package org.example.algorithm.heap;

import java.util.*;

//给你一个整数数组 nums 和一个整数 k ，请你返回其中出现频率前 k 高的元素。你可以按 任意顺序 返回答案。

public class TopKFrequentSolution {

    public static void main(String[] args) {
        int[] nums = {1,1,1,2,2,3};
        int k = 2;
        TopKFrequentSolution solution = new TopKFrequentSolution();
        int[] res = solution.topKFrequent3(nums, k);
        for (int i=0;i<res.length;i++) {
            System.out.println(res[i]);
        }
    }

    //使用快排思想
    public int[] topKFrequent3(int[] nums, int k) {
        Map<Integer, Integer> map = new HashMap<>();
        for (int i = 0; i < nums.length; i++) {
            int count = map.getOrDefault(nums[i], 0);
            map.put(nums[i], count + 1);
        }
        int[][] numCount = new int[map.size()][2];
        int index = 0;
        for (Integer key: map.keySet()) {
            numCount[index++] = new int[]{key, map.get(key)};
        }
        int[] ans = new int[k];
        quickSort(numCount, 0, numCount.length-1, k, ans, 0);
        return ans;
    }

    public void quickSort(int[][] arr, int start, int end, int k, int[] ans, int index) {
        if (start > end || k == 0) {
            return;
        }
        int pos = quickPartition(arr, start, end);
        if (pos - start + 1 <= k) {
            for (int i=start;i<=pos;i++) {
                ans[index++] = arr[i][0];
            }
            quickSort(arr, pos+1, end, k-(pos - start + 1), ans, index);
        } else {
            quickSort(arr, start, pos, k, ans, index);
        }
    }

    private int quickPartition(int[][] arr, int start, int end) {
        Random random = new Random();
        int randomIndex = random.nextInt(end - start + 1) + start;
        swap(arr, start, randomIndex);
        int[] pivot = arr[start];
        int left = start;
        for (int i=start+1;i<=end;i++) {
            if (arr[i][1] >= pivot[1]) {
                left++;
                swap(arr, left, i);
            }
        }
        swap(arr, left, start);
        return left;
    }

    //使用优先级队列
    public int[] topKFrequent2(int[] nums, int k) {
        PriorityQueue<int[]> queue = new PriorityQueue<>((a, b) -> b[1] - a[1]);
        Map<Integer, Integer> map = new HashMap<>();
        for (int i=0;i<nums.length;i++) {
            int count = map.getOrDefault(nums[i], 0);
            map.put(nums[i], count+1);
        }
        for (Integer key: map.keySet()) {
            queue.offer(new int[]{key, map.get(key)});
        }
        int[] ans = new int[k];
        int index = 0;
        while (!queue.isEmpty() && index < k) {
            ans[index++] = queue.poll()[0];
        }
        return ans;
    }

    //构建大顶堆
    public int[] topKFrequent(int[] nums, int k) {
        Map<Integer, Integer> map = new HashMap<>();
        for (int i=0;i<nums.length;i++) {
            int count = map.getOrDefault(nums[i], 0);
            map.put(nums[i], count+1);
        }
        int heapSize = map.size();
        int[][] heap = new int[heapSize][2];
        int index = 0;
        for (Integer key: map.keySet()) {
            heap[index++] = new int[]{key, map.get(key)};
        }
        //构建堆
        int lastNoLeafPos = (heapSize - 2) / 2;
        for (int i=lastNoLeafPos;i>=0;i--) {
            siftDown(heap, i, heapSize);
        }
        int[] ans = new int[k];
        for (int i=0;i<k;i++) {
            ans[i] = removeTop(heap, heapSize - i);
        }
        return ans;
    }

    private int removeTop(int[][] heap, int heapSize) {
        int ans = heap[0][0];
        heap[0] = heap[heapSize-1];
        siftDown(heap, 0, heapSize-1);
        return ans;
    }

    private void siftDown(int[][] heap, int pos, int heapSize) {
        int leftChild = pos * 2 + 1;
        int rightChild = pos * 2 + 2;
        int maxPos = pos;
        if (leftChild <= heapSize-1 && heap[maxPos][1] < heap[leftChild][1]) {
            maxPos = leftChild;
        }
        if (rightChild <= heapSize-1 && heap[maxPos][1] < heap[rightChild][1]) {
            maxPos = rightChild;
        }
        if (maxPos != pos) {
            swap(heap, pos, maxPos);
            siftDown(heap, maxPos, heapSize);
        }
    }

    private void swap(int[][] arr, int from, int to) {
        int[] temp = arr[from];
        arr[from] = arr[to];
        arr[to] = temp;
    }
}
